/*
 * Javlov - a Java toolkit for reinforcement learning with multi-agent support.
 * 
 * Copyright (c) 2009 Matthijs Snel
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

public abstract class TabularQFunction implements QFunction {

	private static final long serialVersionUID = -1678305368586324288L;

	/**
	 * Use this factory method if the size of the state space is unknown. Will result in slower
	 * performance than when using {@link #getInstance(int, int)}.
	 * 
	 * If the size of the state space is unknown because the space is continuous, you probably
	 * shouldn't be using this class, but one of the classes that implement function approximators /
	 * generalisations over states for the value function (yet to be implemented).
	 * 
	 * The use case for this method is one in which the state space is discrete and of 
	 * unknown size (e.g. because the environment is dynamic and hence the size of the state space
	 * is hard to predict), yet finite and small enough for it to make sense to be stored in a
	 * table (as a rough indication I'd say less than a million states).
	 * 
	 * @param actionSpaceSize the (upper bound on the) size of the action space.
	 * @param initValue the initial value for the Q-values, used to initialise the table.
	 */
	public static TabularQFunction getInstance(int actionSpaceSize, double initValue) {
		return new UnknownStateSpaceFunction(actionSpaceSize, initValue);
	}
	
	/**
	 * Use this factory method if the size of the state space is known. Will result in faster
	 * performance than when using {@link #getInstance(int)}.
	 * 
	 * When using this method, the value function constructs an array of the size of the
	 * state space size, and uses the ID of the states to index into the array. Therefore, the
	 * following points of caution should be observed:
	 * 
	 * 1. The provided state space size can not be exceeded (i.e. the actual state space size
	 * can not be larger than the provided one), or errors might start occurring when using 
	 * the value function.
	 * 
	 * 2. States don't change; more particularly, the state ID uniquely identifies a state. For example,
	 * in a static grid world, a state with a particular ID will always be the same (position
	 * in grid, objects near this position, etc).
	 * 
	 * @param stateSpaceSize the (upper bound on the) size of the state space.
	 * @param actionSpaceSize the (upper bound on the) size of the action space.
	 * @param initValue the initial value for the Q-values, used to initialise the table.
	 */
	public static TabularQFunction getInstance(int stateSpaceSize, int actionSpaceSize, 
			double initValue) {
		return new KnownStateSpaceFunction(stateSpaceSize, actionSpaceSize, initValue);
	}
	
	protected double alpha;
	
	/**
	 * Default constructor, sets the learning rate alpha to 0.1
	 */
	protected TabularQFunction() {
		setLearnRate(0.1);
	}
	
	/**
	 * @inheritDoc
	 */
	@Override
	public double getLearnRate() {
		return alpha;
	}

	/**
	 * @inheritDoc
	 */
	@Override
	public void setLearnRate(double alpha) {
		this.alpha = alpha;
	}
	
	protected static class UnknownStateSpaceFunction extends TabularQFunction {

		private static final long serialVersionUID = 8186346862373024852L;

		//Performance test showed Trove only outperforms Java HashMap if state space larger than
		//around 250,000 states or so; for smaller spaces, it is actually slower. On the assumption
		//that most uses will involve less states, using HashMap.
		protected Map<State, double[]> table;
		
		protected int actionSpaceSize;
		
		protected double initVal, initVals[];
		
		private State lastState, prevLastState;
		private double[] lastVals, prevLastVals;
		private int lastActionID, prevLastActionID;
		
		public UnknownStateSpaceFunction(int actionSpaceSize, double initValue) {
			table = new HashMap<State, double[]>(5000);
			this.actionSpaceSize = actionSpaceSize;
			initVal = initValue;
			initVals = getInitValues();
		}
		
		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State, Action)
		 */
		@Override
		public <T> double getValue(State<T> s, Action act) {
			updateMemory(s, act);
			if ( lastVals == null )
				return initVals[lastActionID];
			return lastVals[lastActionID];
		}

		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State)
		 */
		@Override
		public <T> double[] getValues(State<T> s) {
			updateMemory(s, null);
			if ( lastVals == null )
				return initVals;
			return lastVals;
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void setValue(State<T> s, Action act, double value) {
			double[] vals = table.get(s);
			if ( vals == null ) {
				vals = getInitValues();
				table.put(s, vals);
			}
			vals[act.getID()] = value;
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void setLastAction(Action a) {
			lastActionID = a.getID();
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void update(State<T> s, Action act, double delta) {
			update(s, table.get(s), act.getID(), delta);
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void update(double delta) {
			update(lastState, lastVals, lastActionID, delta);
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void updatePrevious(double delta) {
			update(prevLastState, prevLastVals, prevLastActionID, delta);
		}
		
		/**
		 * Initialises the Q-values to the intial values specified in the constructor.
		 * Override if they should be intialised differently.
		 * @return
		 */
		protected double[] getInitValues() {
			double[] vals = new double[actionSpaceSize];
			for ( int i = 0; i < actionSpaceSize; i++ )
				vals[i] = initVal;
			return vals; 
		}
		
		protected <T> void update(State<T> s, double[] vals, int actID, double delta) {
			if ( vals == null ) {
				vals = getInitValues();
				table.put(s, vals);
			}
			vals[actID] += alpha*delta;
		}
		
		/**
		 * Updates the function's "memory" of previous states, actions, and retrieved
		 * (s,a) values, in order to be able to implement methods {@link #update(double)}
		 * and {@link #updatePrevious(double)}.
		 * 
		 * @param s the state currently being queried
		 * @param act the action currently being queried - can be null
		 */
		protected <T> void updateMemory(State<T> s, Action act) {
			prevLastState = lastState;
			lastState = s;
			prevLastVals = lastVals;
			lastVals = table.get(s);
			prevLastActionID = lastActionID;
			if ( act == null )
				lastActionID = 0;
			else
				lastActionID = act.getID();
		}
		
		public void init() {
			reset();
			table.clear();
		}
		
		public void reset() {
			prevLastState = lastState = null;
			prevLastVals = lastVals = null;
			prevLastActionID = lastActionID = 0;
		}
		
		public String toString() {
			String s = "";
			for ( State st : table.keySet() ) {
				s += st.toString() + ":" + Arrays.toString(table.get(st)) + "\n";
			}
			return s;
		}
	}
	
	
	protected static class KnownStateSpaceFunction extends TabularQFunction {
		
		private static final long serialVersionUID = 5529805948086768159L;
		
		/**
		 * Table that stores (s,a) values. 1st index = states, 2nd = actions.
		 */
		protected double[][] table;
		
		private int lastStateID, prevLastStateID;
		private int lastActionID, prevLastActionID;
		
		public KnownStateSpaceFunction(int stateSpaceSize, int actionSpaceSize,
				double initValue) {
			table = getInitTable(stateSpaceSize, actionSpaceSize, initValue);
		}
		
		/**
		 * Initialises the table to 0 values. Override if it should be initialised differently.
		 * @param stateSpaceSize
		 * @param actionSpaceSize
		 * @return
		 */
		protected double[][] getInitTable(int stateSpaceSize, int actionSpaceSize,
				double initValue) {
			double[][] init = new double[stateSpaceSize][actionSpaceSize];
			for ( int i = 0; i < stateSpaceSize; i++ )
				for ( int j = 0; j < actionSpaceSize; j++ )
					init[i][j] = initValue;
			return init;
		}
		
		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State, Action)
		 */
		@Override
		public <T> double getValue(State<T> s, Action act) {
			updateMemory(s, act);
			return table[lastStateID][lastActionID];
		}

		/**
		 *
		 * Returns 0 for (s,a) pairs that haven't been encountered before.
		 * 
		 * @inheritDoc
		 * @see QFunction#getValue(State)
		 */
		@Override
		public <T> double[] getValues(State<T> s) {
			updateMemory(s, null);
			return table[lastStateID];
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void setValue(State<T> s, Action act, double value) {
			table[s.getID()][act.getID()] = value;	
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public void setLastAction(Action a) {
			lastActionID = a.getID();
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void update(State<T> s, Action act, double delta) {
			table[s.getID()][act.getID()] += alpha*delta;			
		}
		
		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void update(double delta) {
			table[lastStateID][lastActionID] += alpha*delta;	
		}

		/**
		 * @inheritDoc
		 */
		@Override
		public <T> void updatePrevious(double delta) {
			table[prevLastStateID][prevLastActionID] += alpha*delta;
		}
		
		/**
		 * Updates the function's "memory" of previous states, actions, and retrieved
		 * (s,a) values, in order to be able to implement methods {@link #update(double)}
		 * and {@link #updatePrevious(double)}.
		 * 
		 * @param s the state currently being queried
		 * @param act the action currently being queried - can be null
		 */
		protected <T> void updateMemory(State<T> s, Action act) {
			prevLastStateID = lastStateID;
			lastStateID = s.getID();
			prevLastActionID = lastActionID;
			if ( act == null )
				lastActionID = 0;
			else
				lastActionID = act.getID();
		}
		
		public void init() {
			reset();
			table = new double[table.length][table[0].length];
		}
		
		public void reset() {
			prevLastStateID = lastStateID = 0;
			prevLastActionID = lastActionID = 0;
		}
	}
}
